/* benchmark guru interfaces */
#include <Utilities/RandomGenerator.h>
#include <Configuration.h>
#include <mpi/communicator.h>
#include <mpi/collectives.h>
#include <Utilities/Timer.h>
#include <OhmmsPETE/OhmmsMatrix.h>
#include <OhmmsPETE/TinyVector.h>
#include <Message/OpenMP.h>
#include <benchmark/transpose.h>
#include <fstream>
#include <mpp/shmem.h>
using namespace APPNAMESPACE;


template<typename T>
inline void check_transpose(Matrix<T>& A,  Matrix<T>& trA, ostream& os)
{
  os << "Before transposes " << endl;
  os << A << endl;
  bool fine=true;
  const double eps=numeric_limits<T>::epsilon()*10.;
  for(int i=0; i<A.rows(); ++i)
    for(int j=0; j<A.cols(); ++j)
      if(abs(A(i,j)-trA(j,i))>eps) fine=false;
  if(fine)
    os<< "Transpose is correct " << endl;
  else
  {
    os << "Incorrect tranpose " << endl;
    os<< trA << endl;
  }
}

template<typename T, unsigned TRANSENG>
struct TransposeHelper
{
  int sizeN;
  int sizeM;
  int locN;
  int locM;
  int block_size;
  int num_blocks;
  TinyVector<double,3> tt;
  const mpi::communicator& mycomm;
  T* mybuffer;
  //double* mybuffer;

  inline TransposeHelper(int n, int m, const mpi::communicator& c)
    :sizeN(n),sizeM(m),tt(0.0),mycomm(c)//,mybuffer(0)
  {
    num_blocks=mycomm.size();
    locN=sizeN/num_blocks;
    locM=sizeM/num_blocks;
    block_size=locN*locM;
    mybuffer=(T*)shmalloc(num_blocks*block_size*sizeof(T));
    //mybuffer=(double*)shmalloc(num_blocks*block_size*sizeof(double));
    //cout << " Checking " << mycomm.size() << " " << mycomm.rank() << " " << num_blocks*block_size << endl;
    //cout << " Checking " << mycomm.size() << " " << mycomm.rank() << endl;
  }

  //destructor
  ~TransposeHelper() 
  {
    if(mybuffer) shfree(mybuffer);
  }

  inline void pack(const T* restrict in, T* restrict out)
  {
    T* restrict tptr=out;
    for(int ip=0; ip<num_blocks; ++ip)
    {
      for(int i=0; i<locM; ++i)
      {
        const T* aptr=in+ip*locM+i;
        for(int j=0; j<locN; ++j,aptr+=sizeM) *tptr++=*aptr;
      }
    }
  }

  inline void unpack(T* restrict out)
  {
    for(int ip=0; ip<num_blocks;++ip)
    {
      //const T* aptr=mybuffer.data()+ip*block_size;
      const T* aptr=mybuffer+ip*block_size;
      for(int i=0,iny=ip; i<locM; ++i,iny+=num_blocks)
      {
        T* tptr=out+iny*locN;
        for(int j=0; j<locN; ++j) *tptr++=*aptr++;
      }
    }
  }

  /** transpose 
   * @param in starting address of the input data(locN,sizeM)
   * @param out starting address of the output data(locM,sizeN)
   *
   * out is temporarily used for alltoall
   */
//  inline void apply(const T* restrict in, T* restrict out)
  inline void apply(const T* in, T* out)
  {
    //int nd=num_blocks*block_size;
    //char fname[64];
    //sprintf(fname,"test.%d",mycomm.rank());
    //ofstream fout(fname);
    //fout << "In " << endl;
    //std::copy(in,in+nd,ostream_iterator<T>(fout," "));
    //fout << endl;

    //pack(in,mybuffer);
    pack(in,out);

    //shmem_udcflush();
    //shmem_barrier_all();
    //pack(in,out);

    //int itarget=(mycomm.rank()+1)%mycomm.size();
    ////cout << mycomm.rank() << " put " << block_size*itarget << " in " << itarget << endl;
    //shmem_double_put(mybuffer+block_size*itarget,out,block_size,itarget);
    
    //for(int i=0; i<num_blocks; ++i)
    int offset_=0;
    int toffset=block_size*mycomm.rank();
    for(int i=0; i<mycomm.size(); ++i,offset_+=block_size)
      shmem_double_put(mybuffer+toffset,out+offset_,block_size,i);
    shmem_barrier_all();
    shmem_udcflush();
    unpack(out);
    ////std::copy(out,out+nd,ostream_iterator<T>(fout," "));
    //std::copy(mybuffer,mybuffer+nd,ostream_iterator<T>(fout," "));
    //fout << endl;
    ////shmem_udcflush();
    ////shmem_barrier_all();

    //fout << "Put " << endl;
    //for(int i=0; i<num_blocks; ++i)
    //{
    //  std::copy(target[i],target[i]+block_size,ostream_iterator<T>(fout," "));
    //  fout << endl;
    //}
    //
    ////shmem_barrier_all();
    //std::copy(mybuffer,mybuffer+block_size*num_blocks,out);
    //unpack(out);
  }

  /** transpose a Matrix
   */
  inline void apply(const Matrix<T>& a, Matrix<T>& trA)
  {
    if(num_blocks>1)
      apply(a.data(),trA.data());
    else
      Transpose2D<T,TRANSENG>::apply(a,trA);
  }
};

int main(int argc, char** argv)
{
  mpi::environment env(argc,argv);
  mpi::communicator mycomm;
  //shmem_init();
  OhmmsInfo ohmms("transpose",mycomm.rank(),0,1);

  typedef  double real_type;
  ///number of iterations
  int niters=10;
  ///rows
  int sizeN=4;
  ///columns
  int sizeM=4;
  ///partition 1
  int nx=1;
  ///partition 2
  int ny=4;
  ///if true, debug the implemnetation
  bool debug=false;

  int ic=0;
  while(ic<argc)
  {
    string a(argv[ic]);
    if(a == "opt_s") 
      sizeN=sizeM=atoi(argv[++ic]);
    else if(a == "opt_n")
      sizeN=atoi(argv[++ic]);
    else if(a == "opt_m")
      sizeM=atoi(argv[++ic]);
    else if(a == "opt_p") 
      ny=atoi(argv[++ic]);
    else if(a == "opt_i")
      niters=atoi(argv[++ic]);
    else if(a == "opt_debug")
      debug=true;
    ++ic;
  }

  ny=mycomm.size();
  int locN=sizeN/ny;
  int locM=sizeM/ny;

  typedef double value_type;

  //typedef int value_type;
  static Matrix<value_type>  A(locN,sizeM), trA(locM,sizeN);
  //TransposeHelper<value_type,0> forward(sizeN,sizeM,mycomm);
  TransposeHelper<value_type,0> *forward = new TransposeHelper<value_type,0>(sizeN,sizeM,mycomm);

  //use A(i,j)=i*sizeM+j  where i=global locator
  for(int i=0; i<locN; ++i)
    for(int j=0; j<sizeM; ++j)
      A(i,j)=(locN*mycomm.rank()+i)*sizeM+j;

  if(debug)
  {
    forward->apply(A,trA);
    bool success=true;
    for(int i=0; i<locM; ++i)
    {
      for(int j=0; j<sizeN; ++j)
        if(abs(j*sizeM+mycomm.rank()*locM+i-trA(i,j))>1e-6) success=false;
    }
    if(success)
      cerr << mycomm.rank() << "  OK" << endl;
    else
      cerr << mycomm.rank() << "  failed" << endl;
  }

  Timer myclock;
  vector<double> tt(4,0.0);
  //for(int i=0; i<niters; ++i)
  //{
  //  for(int i=0; i<A.size(); ++i) A(i)=Random();
  //  myclock.restart();
  //  forward->apply(A,trA);
  //  tt[0]+=myclock.elapsed();
  //}
  //for(int i=0; i<3; ++i) tt[i+1]=forward->tt[i];
  //mpi::reduce(mycomm,tt);

  app_log() << "Transpose " << sizeN << " " << sizeM << " "
    << locN << " " << locM << " " 
    << locN*locM*sizeof(value_type) << " "
    << tt[0]/(mycomm.size()*niters) << " "
    << tt[2]/(mycomm.size()*niters) << " "
    << tt[1]/(mycomm.size()*niters) << " "
    << tt[3]/(mycomm.size()*niters) 
    << endl;

  delete forward;
  //shmem_finalize();
  return 0;
}

